from glob import glob
import json
import cv2
import numpy as np

"""
输入放着原始图像和标注数据的路径，
1. 通过该路径获取所有图像和json文件列表，
2. 提取所有json文件中的标签，放到分类列表中，
3. 根据分类列表中的索引，生成mask,
4. mask是灰度图，背景像素值为0，类别的像素值依次为1、2、3...
"""


def get_file_list(path):
    # 根据路径获取其中的jpg格式原图像，和json格式标注文件，放到字典中返回
    base_path = path
    images = glob(base_path + '\\*.jpg')
    jsonfiles = glob(base_path + '\\*.json')
    return {"images": images, "jsonfiles": jsonfiles}


def category_list_generate(jsonfiles):
    # 根据json文件列表，获取所有种类的标签，放到一个列表中返回
    category_types = ["Background"]
    for file in jsonfiles:
        with open(file, "r") as f:
            data = json.load(f)
        shapes = data["shapes"]
        for shape in shapes:
            label = shape["label"]
            cls_name = label
            if cls_name not in category_types:
                category_types.append(cls_name)
    return category_types


def json_to_dataset(category_types, image, jsonfile, out_path):
    # 根据一个存放着标签的列表，以及原始图像文件路径和json标注文件路径，生成8位灰度图像，输出到out_path
    filename = image.split("\\")[-1].split(".")[0]
    img = cv2.imread(image)
    h, w = img.shape[:2]
    mask = np.zeros([h, w, 1], np.uint8)

    with open(jsonfile, "r") as f:
        label = json.load(f)

    shapes = label["shapes"]
    for shape in shapes:
        category = shape["label"]
        points = shape["points"]
        points_array = np.array(points, dtype=np.int32)
        mask = cv2.fillPoly(mask, [points_array], category_types.index(category))

    cv2.imwrite(out_path + '\\' + filename + ".png", mask)


def main():
    dir_path = 'D:\\Files\\deepLearningData\\lcx_lad1'  # 文件所在路径
    annotations_path = dir_path + "\\annotations"
    out_category_types_file = annotations_path + "\\class_name.txt"
    dic = get_file_list(dir_path)
    images = dic["images"]
    jsonfiles = dic["jsonfiles"]
    category_types = category_list_generate(jsonfiles)
    with open(out_category_types_file, 'w') as f:
        f.writelines('\n'.join(category_types))  # writelines可将列表中的字符串写入文本
    for i in range(len(images)):
        json_to_dataset(category_types, images[i], jsonfiles[i], annotations_path)


if __name__ == '__main__':
    main()
